'''
@author: kohill
@Date: 2018-8-7
This file & resnet_v1_101_rfcn_dcn.py is original RFCN model but rewritten with Gluon, as
the original implementation written by Symbol API is not flexible enough.
For pre-trained model and other information, you can take a look at <https://github.com/msracver/Deformable-ConvNets>
and the paper <https://arxiv.org/abs/1703.06211>.
'''
from __future__ import print_function
import mxnet as mx
import mxnet.gluon.nn as nn
import mxnet.ndarray as nd
import numpy as np
import mxnet.autograd as ag
from models.resnet_v1_101_rfcn_dcn import resnet_v1_101_rfcn_dcn as sym_res101
from models.proposal import Proposal,Proposal_Target
from lib.mxnet.params import load_param
from lib.bbox.bbox_transform import bbox_pred,clip_boxes
import gluoncv
import matplotlib.pyplot as plt

class Resnet(nn.HybridBlock):
    def __init__(self,cfg):
        super(Resnet,self).__init__()
        Sym = sym_res101()
        inp = mx.sym.Variable("data")
        sym_feat4 = Sym.get_resnet_v1_conv4(data = inp)
        sym_feat5 = Sym.get_resnet_v1_conv5(sym_feat4)
        sym = mx.sym.Group([sym_feat4,sym_feat5])
        with self.name_scope():
            self.features = nn.SymbolBlock(sym, inp)
        if cfg.pretrained is not None:
            '''
            load parameters trained from Deformable-ConvNets, or the pretrained model trained on ImageNet.
            '''
            args,auxes = load_param(cfg.pretrained)
            net_params = self.features.collect_params()
            ctx = mx.cpu()
            net_params["res5a_branch2b_offset_weight"].initialize(init = mx.init.Zero())
            net_params["res5a_branch2b_offset_bias"].initialize(init = mx.init.Zero())
            net_params["res5b_branch2b_offset_weight"].initialize(init = mx.init.Zero())
            net_params["res5b_branch2b_offset_bias"].initialize(init = mx.init.Zero())
            net_params["res5c_branch2b_offset_weight"].initialize(init = mx.init.Zero())
            net_params["res5c_branch2b_offset_bias"].initialize(init = mx.init.Zero())

            for param in args.keys():
                if param in net_params:
                    net_params[param]._load_init(args[param], ctx=ctx)
            for param in auxes.keys():
                if param in net_params:
                    net_params[param]._load_init(auxes[param], ctx=ctx)
        if cfg.resume is not None:
            args = mx.nd.load(cfg.resume)
            for param in args.keys():
                net_params = self.features.collect_params()
                if param in net_params:
                    net_params[param]._load_init(args[param], ctx=mx.cpu())

    def forward(self, x):

        return self.features(x)

class RFCN_Resnet(mx.gluon.nn.HybridBlock):
    def __init__(self,cfg,is_train = False,pretrained = True,viz_env = None):
        super(RFCN_Resnet,self).__init__()
        self.cfg = cfg
        self._is_train = is_train

        num_classes = cfg.dataset.NUM_CLASSES
        num_reg_classes = 2
        self.num_classes = num_classes
        self.num_reg_classes = num_reg_classes
        self.feat = Resnet(self.cfg)
        # self.feat.collect_params().setattr('lr_mult', 0)
        self.viz_env = None
        if viz_env is not None:
            from visdom import Visdom
            self.viz =  Visdom(port = 8098)
        with self.name_scope():
            self.rpn_relu = nn.Conv2D(channels=512,kernel_size=(3,3),strides=(1,1),padding=(1,1),activation="relu")

            self.rpn_cls_score = nn.Conv2D(channels=2 * cfg.network.NUM_ANCHORS,kernel_size=(1,1),strides=(1,1))
            self.relu_new_1 = nn.Conv2D(channels=1024,kernel_size=(1,1),strides=(1,1),activation="relu")
            self.relu_new_1.collect_params().setattr('lr_mult', 3)
            self.rpn_bbox_pred = nn.Conv2D(channels=4 * cfg.network.NUM_ANCHORS,kernel_size=(1,1),strides=(1,1))
            self.rfcn_cls = nn.Conv2D(channels=7 * 7 * num_classes, kernel_size=(1, 1), strides=(1, 1))
            self.rfcn_bbox = nn.Conv2D(channels=7 * 7 * num_reg_classes * 4, kernel_size=(1, 1), strides=(1, 1))
            self.rfcn_cls_offset_t = nn.Conv2D(channels=2 * 7 * 7 * num_classes,kernel_size=(1,1),strides=(1,1))
            self.rfcn_bbox_offset_t = nn.Conv2D(channels=2 * 7 * 7,kernel_size=(1,1),strides=(1,1))

            if is_train:self.proposal_target = Proposal_Target(cfg)
            self.proposal = Proposal(cfg)
        if self.cfg.pretrained and pretrained:
            args, auxes = load_param(cfg.pretrained)
            self.rfcn_bbox_offset_t.weight._load_init(args["rfcn_bbox_offset_t_weight"], ctx=mx.cpu())
            self.rfcn_bbox_offset_t.bias._load_init(args["rfcn_bbox_offset_t_bias"], ctx=mx.cpu())
            self.rfcn_cls_offset_t.weight._load_init(args["rfcn_cls_offset_t_weight"], ctx=mx.cpu())
            self.rfcn_cls_offset_t.bias._load_init(args["rfcn_cls_offset_t_bias"], ctx=mx.cpu())
            self.rfcn_bbox.weight._load_init(args["rfcn_bbox_weight"], ctx=mx.cpu())
            self.rfcn_bbox.bias._load_init(args["rfcn_bbox_bias"], ctx=mx.cpu())
            self.rfcn_cls.weight._load_init(args["rfcn_cls_weight"], ctx=mx.cpu())
            self.rfcn_cls.bias._load_init(args["rfcn_cls_bias"], ctx=mx.cpu())
            self.rfcn_cls.collect_params().setattr('lr_mult', 1)
            self.relu_new_1.weight._load_init(args["conv_new_1_weight"], ctx=mx.cpu())
            self.relu_new_1.bias._load_init(args["conv_new_1_bias"], ctx=mx.cpu())

            self.rpn_relu.weight._load_init(args["rpn_conv_3x3_weight"], ctx=mx.cpu())
            self.rpn_relu.bias._load_init(args["rpn_conv_3x3_bias"], ctx=mx.cpu())

            self.rpn_bbox_pred.weight._load_init(args["rpn_bbox_pred_weight"], ctx=mx.cpu())
            self.rpn_bbox_pred.bias._load_init(args["rpn_bbox_pred_bias"], ctx=mx.cpu())

            self.rpn_cls_score.weight._load_init(args["rpn_cls_score_weight"], ctx=mx.cpu())
            self.rpn_cls_score.bias._load_init(args["rpn_cls_score_bias"], ctx=mx.cpu())
        else:
            self.rfcn_bbox_offset_t.weight.initialize(init=mx.init.Zero())
            self.rfcn_bbox_offset_t.bias.initialize(init=mx.init.Zero())
            self.rfcn_cls_offset_t.weight.initialize(init=mx.init.Zero())
            self.rfcn_cls_offset_t.bias.initialize(init=mx.init.Zero())
            self.rfcn_bbox.weight.initialize(init=mx.init.Normal(sigma=0.01), ctx=mx.cpu())
            self.rfcn_bbox.bias.initialize(init=mx.init.Zero(), ctx=mx.cpu())
            self.rfcn_cls.weight.initialize(init=mx.init.Normal(sigma=0.01), ctx=mx.cpu())
            self.rfcn_cls.bias.initialize(init=mx.init.Zero(), ctx=mx.cpu())
            self.relu_new_1.weight.initialize(init=mx.init.Normal(sigma=0.01), ctx=mx.cpu())
            self.relu_new_1.bias.initialize(init=mx.init.Zero(), ctx=mx.cpu())

            self.rpn_relu.weight.initialize(init=mx.init.Normal(sigma=0.01), ctx=mx.cpu())
            self.rpn_relu.bias.initialize(init=mx.init.Zero(), ctx=mx.cpu())

            self.rpn_bbox_pred.weight.initialize(init=mx.init.Normal(sigma=0.01), ctx=mx.cpu())
            self.rpn_bbox_pred.bias.initialize(init=mx.init.Zero(), ctx=mx.cpu())

            self.rpn_cls_score.weight.initialize(init=mx.init.Normal(sigma=0.01), ctx=mx.cpu())
            self.rpn_cls_score.bias.initialize(init=mx.init.Zero(), ctx=mx.cpu())
        if cfg.resume is not None:
            self.collect_params().load(cfg.resume,allow_missing = True)
    def rcnn_calc(self,batch_ori_img,relu1,rpn_cls_act_reshape,rpn_bbox_pred,im_info,gt_boxes,num_classes,num_reg_classes):

        F = mx.nd
        cfg = self.cfg

        if ag.is_training():
        
            rois, rcnn_label,rcnn_label_onehot, rcnn_bbox_target, rcnn_bbox_weight = self.proposal_target(rpn_cls_act_reshape, rpn_bbox_pred, im_info ,gt_boxes[:,:,:5].asnumpy())

            if self.cfg.TRAIN.VIZ_RPN:
                '''Visualize rcnn label_to make sure it's correct.'''
                assert np.all(rcnn_label != -1)
                roi_index = np.where(rois[:,0]==0)[0]
                roi_index_label1 = np.where(rcnn_label[roi_index]>0)[0]
                t_rois = rois[roi_index,1:]
                gluoncv.utils.viz.plot_bbox(batch_ori_img[0],t_rois[roi_index_label1],labels=rcnn_label[roi_index][roi_index_label1])
                self.viz.matplot(plt,env=self.viz_env,win="rpn_rois")
                plt.close()
                '''Visualize rcnn_bbox_target to make sure it's correct, terribly I found there is something wrong.'''

                rois = mx.nd.array(rois,ctx = relu1.context)
                rcnn_label = mx.nd.array(rcnn_label,ctx = relu1.context)
                rcnn_label_onehot = mx.nd.array(rcnn_label_onehot,ctx = relu1.context)
                rcnn_bbox_target = mx.nd.array(rcnn_bbox_target,ctx = relu1.context)
                rcnn_bbox_weight = mx.nd.array(rcnn_bbox_weight,ctx = relu1.context)
                from lib.nms.nms import gpu_nms_wrapper
                _pred_bboxes, _pred_scores, _pred_labels = self.post_process(output_rois = rois[roi_index][roi_index_label1],
                                                                             output_bbox_pred = rcnn_bbox_target[roi_index][roi_index_label1].reshape(1,-1,8),
                                                                             output_cls_score = rcnn_label_onehot[roi_index][roi_index_label1].reshape(1,-1,self.num_classes) ,
                                                                             nms_wrapper = gpu_nms_wrapper(0.7,0),
                                                                             im_shape = [batch_ori_img[0].shape[2],batch_ori_img[0].shape[0],batch_ori_img[0].shape[1]],
                                                                             scale = 1.0)
                gluoncv.utils.viz.plot_bbox(batch_ori_img[0],
                                            bboxes= _pred_bboxes,
                                            labels= _pred_labels,
                                            scores= _pred_scores
                                            )
                self.viz.matplot(plt,env=self.viz_env,win="rpn_rois_target")
                plt.close()
            else:
                rois = mx.nd.array(rois, ctx=relu1.context)
                rcnn_label = mx.nd.array(rcnn_label, ctx=relu1.context)
                rcnn_label_onehot = mx.nd.array(rcnn_label_onehot, ctx=relu1.context)
                rcnn_bbox_target = mx.nd.array(rcnn_bbox_target, ctx=relu1.context)
                rcnn_bbox_weight = mx.nd.array(rcnn_bbox_weight, ctx=relu1.context)
        else:
            rpn_rois, rpn_scores = self.proposal(rpn_cls_act_reshape.asnumpy(), rpn_bbox_pred.asnumpy(), im_info.asnumpy())
            rois = rpn_rois.reshape((-1, rpn_rois.shape[2]))
            scores = rpn_scores.reshape((-1, rpn_scores.shape[2]))
            roi_index = np.where(rois[:,0]==0)[0]
            if self.viz_env is not None:
                t_rois = rois[roi_index,1:]
                gluoncv.utils.viz.plot_bbox(batch_ori_img[0],t_rois,scores=scores[roi_index],thresh=.8)
                self.viz.matplot(plt,env=self.viz_env,win="rpn_rois_predict")
                plt.close()
            rois = mx.nd.array(rois,ctx=relu1.context)
            rpn_scores = mx.nd.array(rpn_scores, ctx=relu1.context)
        relu_new_1 = self.relu_new_1(relu1)
        rfcn_cls = self.rfcn_cls(relu_new_1)
        rfcn_bbox = self.rfcn_bbox(relu_new_1)



        rfcn_cls_offset_t = self.rfcn_cls_offset_t(relu_new_1)
        rfcn_bbox_offset_t = self.rfcn_bbox_offset_t(relu_new_1)

        rfcn_cls_offset = F.contrib.DeformablePSROIPooling(data=rfcn_cls_offset_t, rois=rois, group_size=7, pooled_size=7,
                                                                sample_per_part=4, no_trans=True, part_size=7, output_dim=2 * num_classes, spatial_scale=1 / self.cfg.network.RPN_FEAT_STRIDE)
        rfcn_bbox_offset = F.contrib.DeformablePSROIPooling(data=rfcn_bbox_offset_t, rois=rois, group_size=7, pooled_size=7,
                                                                 sample_per_part=4, no_trans=True, part_size=7, output_dim=2, spatial_scale=1 / self.cfg.network.RPN_FEAT_STRIDE)
        psroipooled_cls_rois = F.contrib.DeformablePSROIPooling(data=rfcn_cls, rois=rois, trans=rfcn_cls_offset,
                                                                     group_size=7, pooled_size=7, sample_per_part=4, no_trans=False, trans_std=0.1,
                                                                     output_dim=num_classes, spatial_scale=1 / self.cfg.network.RPN_FEAT_STRIDE, part_size=7)
        psroipooled_loc_rois = F.contrib.DeformablePSROIPooling( data=rfcn_bbox, rois=rois, trans=rfcn_bbox_offset,
                                                                     group_size=7, pooled_size=7, sample_per_part=4, no_trans=False, trans_std=0.1,
                                                                     output_dim=8, spatial_scale=1 / self.cfg.network.RPN_FEAT_STRIDE, part_size=7)
        cls_score = F.Pooling(data=psroipooled_cls_rois, pool_type='avg', global_pool=True, kernel=(7, 7))
        bbox_pred = F.Pooling(data=psroipooled_loc_rois, pool_type='avg', global_pool=True, kernel=(7, 7))
        cls_score = F.Reshape(data=cls_score, shape=(-1, num_classes))
        bbox_pred = F.Reshape(data=bbox_pred, shape=(-1, 4 * num_reg_classes))
        if ag.is_training():
            if self.cfg.TRAIN.ENABLE_OHEM:
                from .proposal import box_annotator_ohem
                # labels_ohem, labels_ohem_onehot, labels_ohem_weight, bbox_weights_ohem, label_cnt = box_annotator_ohem(cls_score=cls_score,
                #                                                     bbox_pred=bbox_pred,
                #                                                     labels=rcnn_label,
                #                                                     bbox_targets=rcnn_bbox_target,
                #                                                     bbox_weights=rcnn_bbox_weight,
                #                                                     num_classes = cfg.dataset.NUM_CLASSES,
                #                                                     num_reg_classes = 2 if cfg.CLASS_AGNOSTIC else cfg.dataset.NUM_CLASSES,
                #                                                     roi_per_img=cfg.TRAIN.BATCH_ROIS_OHEM * relu1.shape[0])
                # loss_cls = nd.sum(-1* labels_ohem_onehot * nd.log_softmax(data=cls_score,axis=1) * labels_ohem_weight) / label_cnt
                # loss_l1 = nd.sum(bbox_weights_ohem * nd.smooth_l1(bbox_pred - rcnn_bbox_target ,scalar=3.0)) / label_cnt
                # return loss_cls, loss_l1
                weight_index = box_annotator_ohem(cls_score=cls_score,
                                                                    bbox_pred=bbox_pred,
                                                                    labels=rcnn_label,
                                                                    bbox_targets=rcnn_bbox_target,
                                                                    bbox_weights=rcnn_bbox_weight,
                                                                    num_classes = cfg.dataset.NUM_CLASSES,
                                                                    num_reg_classes = 2 if cfg.CLASS_AGNOSTIC else cfg.dataset.NUM_CLASSES,
                                                                    roi_per_img=cfg.TRAIN.BATCH_ROIS_OHEM * relu1.shape[0])
                cls_prob = mx.nd.log_softmax(data=cls_score,axis=1)
                loss_l1 = rcnn_bbox_weight[weight_index] * mx.nd.smooth_l1(scalar=1.0,
                                                                           data=(bbox_pred[weight_index] - rcnn_bbox_target[weight_index])) / len(weight_index)
                loss_cls = -1 * (cls_prob[weight_index] * rcnn_label_onehot[weight_index]) / len(weight_index)
                return loss_cls, loss_l1
            else:
                cls_prob = mx.nd.log_softmax(data=cls_score,axis=1)
                loss_l1 = rcnn_bbox_weight * mx.nd.smooth_l1(scalar=1.0, data=(bbox_pred - rcnn_bbox_target)) / rois.shape[0]
                loss_cls = -1 * (cls_prob * rcnn_label_onehot) / rois.shape[0]
                return loss_cls, loss_l1
        cls_prob = F.softmax(data=cls_score, axis=1)
        cls_prob = F.Reshape(data=cls_prob, shape=(relu1.shape[0], -1, num_classes))

        return rois,cls_prob,F.Reshape(data=bbox_pred, shape=(relu1.shape[0], -1, 4 * num_reg_classes))
    def forward(self, batch_ori_img,
                x, im_info,
                rpn_cls_label,
                rpn_cls_weight,
                rpn_bbox_target,
                rpn_bbox_weight,
                gt_boxes
                ):
        label = {}
        label["rpn_cls_label"] = rpn_cls_label
        label['rpn_cls_weight'] = rpn_cls_weight
        label["rpn_bbox_target"] = rpn_bbox_target
        label["rpn_bbox_weight"] = rpn_bbox_weight
        label["gt_boxes"] = gt_boxes

        F = nd
        # im_info = nd.array(im_info)
        num_classes = self.num_classes
        num_reg_classes = self.num_reg_classes

        cfg = self.cfg
        num_anchors = cfg.network.NUM_ANCHORS

        conv_feat,relu1 = self.feat(x)
        rpn_relu = self.rpn_relu(conv_feat)
        rpn_cls_score = self.rpn_cls_score(rpn_relu)
        rpn_bbox_pred = self.rpn_bbox_pred(rpn_relu)

        if ag.is_training():
            rpn_cls_score_reshape = F.reshape(rpn_cls_score,(0, 2, -1,rpn_cls_score.shape[2], rpn_cls_score.shape[3]))
            loss_rpn_cls = mx.nd.sum(mx.nd.log_softmax(rpn_cls_score_reshape, axis=1) * label["rpn_cls_label"] *
                                     label["rpn_cls_weight"] * -1) / mx.nd.sum(label["rpn_cls_weight"])

            loss_rpn_l1 = mx.nd.sum(label["rpn_bbox_weight"] *
                                    mx.nd.smooth_l1(scalar=3.0, data=(rpn_bbox_pred - label["rpn_bbox_target"]))) / cfg.TRAIN.RPN_BATCH_SIZE

        rpn_cls_prob = F.softmax(rpn_cls_score.reshape((0,2,-1,0)),axis=1)
        rpn_cls_act_reshape = F.reshape(rpn_cls_prob,(0, 2 * num_anchors, -1, 0))

        if ag.is_training():
            rcnn_loss_cls, rcnn_loss_l1 = \
            self.rcnn_calc(batch_ori_img,relu1,
                       rpn_cls_act_reshape,
                       rpn_bbox_pred,
                       im_info,
                       label["gt_boxes"],
                       num_classes, 
                       num_reg_classes)
            return loss_rpn_cls,loss_rpn_l1,rcnn_loss_cls,rcnn_loss_l1
        else:
            rois,cls_prob,bbox_pred = self.rcnn_calc(batch_ori_img,relu1, rpn_cls_act_reshape, rpn_bbox_pred, im_info, gt_boxes=None, num_classes=num_classes, num_reg_classes=num_reg_classes)
            return rois,cls_prob,bbox_pred
    def post_process(self,output_rois, output_bbox_pred, output_cls_score,nms_wrapper,im_shape,scale):

        scores_all = []
        pred_boxes_all = []

        rois = output_rois.asnumpy()[:,1:]

        # save output
        scores = output_cls_score.asnumpy()[0]
        bbox_deltas = output_bbox_pred.asnumpy()[0]

        if self.cfg.TRAIN.BBOX_NORMALIZATION_PRECOMPUTED:
            BBOX_STDS = np.array(self.cfg.TRAIN.BBOX_STDS)[np.newaxis, np.newaxis, :]
            BBOX_STDS = np.repeat(BBOX_STDS, axis=1, repeats=bbox_deltas.shape[1] // 4)
            BBOX_STDS = np.repeat(BBOX_STDS, axis=0, repeats=bbox_deltas.shape[0])
            bbox_deltas = bbox_deltas * BBOX_STDS.reshape(bbox_deltas.shape)
        pred_boxes = bbox_pred(rois, bbox_deltas)
        pred_boxes = clip_boxes(pred_boxes, im_shape[-2:])

        # we used scaled image & roi to train, so it is necessary to transform them back
        pred_boxes = pred_boxes / scale

        scores_all.append(scores)
        pred_boxes_all.append(pred_boxes)

        scores = scores_all
        boxes = pred_boxes_all

        boxes = boxes[0].astype('f')
        scores = scores[0].astype('f')
        dets_nms = []
        for j in range(1, scores.shape[1]):
            cls_scores = scores[:, j, np.newaxis]
            cls_boxes = boxes[:, 4:8]
            cls_dets = np.hstack((cls_boxes, cls_scores))
            #print(cls_scores)
            keep = nms_wrapper(cls_dets)
            cls_dets = cls_dets[keep, :]
            cls_dets = cls_dets[cls_dets[:, -1] > self.cfg.TEST.KEEP_THRESHOLD , :]
            dets_nms.append(cls_dets)
        
        pred_bboxes = []
        pred_labels = [] 
        pred_scores = []
        for cls_id,det_nms in enumerate(dets_nms):
            if det_nms.shape[0] > 0:
                pred_bboxes.append(det_nms[:,:4])
                pred_scores.append(det_nms[:,4])
                label = np.ones(shape = (det_nms.shape[0],),dtype = "int32")*int(cls_id)
                pred_labels.append(label)
        if len(pred_bboxes) >0 :
            pred_bboxes = np.concatenate(pred_bboxes,axis = 0)
            pred_scores = np.concatenate(pred_scores,axis = 0)
            pred_labels = np.concatenate(pred_labels,axis = 0)
        # visualize
        return np.array(pred_bboxes),np.array(pred_scores),np.array(pred_labels)

